"""
mineru
"""
from typing import Tuple
from PIL import Image
from marker.schema import BlockTypes
from marker.schema.blocks import Block
from typing import Annotated, List
from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm import tqdm
from marker.logger import get_logger

from ..processors import BaseOptProcessor
from ..vlm.vlm_analyze import VlmAnalyzer
from ..config import settings

logger = get_logger()


class OptTableProcessor(BaseOptProcessor):
    block_types: Annotated[
        Tuple[BlockTypes],
        "The block types to process.",
    ] = (BlockTypes.Table, )
    vlm = VlmAnalyzer(settings.VLM_SERVER_URL)
    prompts = ''
    # prompts = 'Table Parsing:'
    # prompts = 'Here is a table, Table Parsing:'

    def process_rewriting(self,
            doc_blocks: List[Block],
            block_images: List[Image.Image]
    ):
        pbar = tqdm(
            desc=f"{self.__class__.__name__} running", disable=self.disable_tqdm
        )
        block_datas = []
        with ThreadPoolExecutor(max_workers=self.max_concurrency) as executor:
            for future in as_completed(
                    [
                        executor.submit(self.single_process, block, block_image)
                        for block, block_image in zip(doc_blocks, block_images)
                    ]
            ):
                if future.exception() is None:
                    block_datas.append(future.result())
                else:
                    logger.error(future.exception())
                pbar.update(1)
        pbar.close()

        return block_datas

    def single_process(self, block: Block, block_image: Image.Image):
        page_markdown = self.vlm.batch_analyze([block_image], self.prompts)[0]
        block_data = {
            'id': str(block.id),
            'html': page_markdown
        }
        return block_data